import taichi as ti

ti.init(arch=ti.cuda)

n = 1080
pixels = ti.Vector.field(3,ti.f32, shape=(n * 2, n))
 
@ti.func
def complex_sqr(z): 
    return ti.Vector([z[0]**2 - z[1]**2,  z[1] * z[0] * 2])


@ti.kernel
def paint(t: float):
    for i, j in pixels:
        c = ti.Vector([-0.8,  ti.cos(t) * 0.2])
        z = ti.Vector([i / n - 1, j / n - 0.5]) * 3
        iterations = 0
        while z.norm() < 20 and iterations < 50:
            z = complex_sqr(z) + c
            iterations += 1
        pixels[i, j] =  [
            (1 - iterations * 0.01),
            (1 - iterations * 0.05),
            (1 - iterations * 0.03)
          ]


gui = ti.GUI("julia set ",res=(n*2,n) )

for i in range(1000000):
    paint(i * 0.03)
    gui.set_image(pixels) 
    gui.show()
